import math
import paddle
from paddle import nn
import paddle.nn.functional as F
from paddle import ParamAttr
from paddle.vision.ops import DeformConv2D


def get_bias_attr(k):
    stdv = 1.0 / math.sqrt(k * 1.0)
    initializer = paddle.nn.initializer.Uniform(-stdv, stdv)
    bias_attr = ParamAttr(initializer=initializer)
    return bias_attr


class Head(nn.Layer):
    def __init__(self, in_channels, dcn=False):
        super(Head, self).__init__()
        self.conv1 = nn.Conv2D(
            in_channels=in_channels,
            out_channels=in_channels // 4,
            kernel_size=3,
            padding=1,
            bias_attr=False)
        self.conv_bn1 = nn.BatchNorm(
            num_channels=in_channels // 4,
            param_attr=ParamAttr(
                initializer=paddle.nn.initializer.Constant(value=1.0)),
            bias_attr=ParamAttr(
                initializer=paddle.nn.initializer.Constant(value=1e-4)),
            act='relu')
        self.conv2 = nn.Conv2DTranspose(
            in_channels=in_channels // 4,
            out_channels=in_channels // 4,
            kernel_size=2,
            stride=2,
            weight_attr=ParamAttr(
                initializer=paddle.nn.initializer.KaimingUniform()),
            bias_attr=get_bias_attr(in_channels // 4))
        self.conv_bn2 = nn.BatchNorm(
            num_channels=in_channels // 4,
            param_attr=ParamAttr(
                initializer=paddle.nn.initializer.Constant(value=1.0)),
            bias_attr=ParamAttr(
                initializer=paddle.nn.initializer.Constant(value=1e-4)),
            act="relu")
        self.conv3 = nn.Conv2DTranspose(
            in_channels=in_channels // 4,
            out_channels=in_channels // 4,
            kernel_size=2,
            stride=2,
            weight_attr=ParamAttr(
                initializer=paddle.nn.initializer.KaimingUniform()),
            bias_attr=get_bias_attr(in_channels // 4),
        )
        self.dcn = dcn

        planes = in_channels // 4

        if dcn:
            self.mask_conv = nn.Conv2D(planes, 9, 3, 1, 1)
            self.offset_conv = nn.Conv2D(planes, 18, 3, 1, 1)      
            self.deform_conv = DeformConv2D(planes, 1, 3, 1, 1)
        else:
            self.out_conv = nn.Conv2D(planes, 1, 3, 1, 1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv_bn1(x)
        x = self.conv2(x)
        x = self.conv_bn2(x)
        x = self.conv3(x)

        if self.dcn:
            mask = F.sigmoid(self.mask_conv(x))
            offset = self.offset_conv(x)
            out = self.deform_conv(x, offset, mask)
        else:
            out = self.out_conv(x)
        out = F.sigmoid(out)
        return out


class DBHeadDeform(nn.Layer):
    """
    Differentiable Binarization (DB) for text detection:
        see https://arxiv.org/abs/1911.08947
    args:
        params(dict): super parameters for build DB network
    """

    def __init__(self, in_channels, k=50, **kwargs):
        super().__init__()
        self.k = k
        binarize_name_list = [
            'conv2d_56', 'batch_norm_47', 'conv2d_transpose_0', 'batch_norm_48',
            'conv2d_transpose_1', 'binarize'
        ]
        thresh_name_list = [
            'conv2d_57', 'batch_norm_49', 'conv2d_transpose_2', 'batch_norm_50',
            'conv2d_transpose_3', 'thresh'
        ]
        self.binarize = Head(in_channels, dcn=True)
        self.thresh = Head(in_channels, dcn=False)

    def step_function(self, x, y):
        return paddle.reciprocal(1 + paddle.exp(-self.k * (x - y)))

    def forward(self, x):
        shrink_maps = self.binarize(x)
        if not self.training:
            return {'maps': shrink_maps}

        threshold_maps = self.thresh(x)
        binary_maps = self.step_function(shrink_maps, threshold_maps)
        y = paddle.concat([shrink_maps, threshold_maps, binary_maps], axis=1)

        return {'maps': y}



class DeformConv(nn.Layer):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1):
        super().__init__()
        pad = kernel_size // 2
        k2 = kernel_size*kernel_size
        self.mask_conv = nn.Conv2D(in_channels, k2, kernel_size, stride, pad)
        self.offset_conv = nn.Conv2D(in_channels, 2*k2, kernel_size, stride, pad)      
        self.deform_conv = DeformConv2D(in_channels, out_channels, kernel_size, stride, pad)

    def forward(self, x):
        mask = F.sigmoid(self.mask_conv(x))
        offset = self.offset_conv(x)
        out = self.deform_conv(x, offset, mask)
        return out


class HeadV2(nn.Layer):
    def __init__(self, in_channels, hidden_dim=64,  dcn=False):
        super().__init__()
        self.up1 = self._make_upsample(in_channels, hidden_dim)
        self.conv1 = nn.Sequential(
            DeformConv(hidden_dim, hidden_dim, 3, 1) if dcn else \
            nn.Conv2D(hidden_dim, hidden_dim, 3, 1, 1, bias_attr=False),
            nn.BatchNorm2D(hidden_dim),
            nn.ReLU()
        )
        self.up2 = self._make_upsample(hidden_dim, hidden_dim)
        self.conv2 = DeformConv(hidden_dim, 1, 3, 1) if dcn else \
            nn.Conv2D(hidden_dim, 1, 3, 1, 1)


    def _make_upsample(self, in_channels, out_channels, scale=2.0):
        return nn.Sequential(
            nn.UpsamplingBilinear2D(scale_factor=scale),
            nn.Conv2D(in_channels, out_channels, 3, 1, 1, bias_attr=False),
            nn.BatchNorm2D(out_channels),
            nn.ReLU()
        )


    def forward(self, x):
        x = self.up1(x)
        x = self.conv1(x)
        x = self.up2(x)
        x = self.conv2(x)
        out = F.sigmoid(x)
        return out



class DBHeadDeformV2(nn.Layer):

    def __init__(self, in_channels, k=50, **kwargs):
        super().__init__()
        self.k = k
        self.binarize = HeadV2(in_channels, dcn=True)
        self.thresh = HeadV2(in_channels, dcn=False)

    def step_function(self, x, y):
        return paddle.reciprocal(1 + paddle.exp(-self.k * (x - y)))

    def forward(self, x):
        shrink_maps = self.binarize(x)
        if not self.training:
            return {'maps': shrink_maps}

        threshold_maps = self.thresh(x)
        binary_maps = self.step_function(shrink_maps, threshold_maps)
        y = paddle.concat([shrink_maps, threshold_maps, binary_maps], axis=1)

        return {'maps': y}